import numpy as np
import cv2
import random
from albumentations.core.transforms_interface import DualTransform
from albumentations.augmentations import functional as F

class MosaicAugmentation(DualTransform):
    """马赛克增强实现，将4张图像拼接成1张"""
    def __init__(self, probability=1.0, img_size=640):
        super().__init__(always_apply=False, p=probability)
        self.img_size = img_size
        self.center_x = int(img_size / 2)
        self.center_y = int(img_size / 2)

    def apply(self, img, **params):
        # 实现马赛克增强逻辑
        h, w = img.shape[:2]
        new_img = np.zeros((self.img_size, self.img_size, 3), dtype=np.uint8)

        # 随机生成马赛克区域大小
        xmin = max(0, int(random.uniform(0, self.center_x)))
        ymin = max(0, int(random.uniform(0, self.center_y)))
        xmax = min(self.img_size, int(random.uniform(self.center_x, self.img_size)))
        ymax = min(self.img_size, int(random.uniform(self.center_y, self.img_size)))

        # 放置原始图像到马赛克区域
        new_img[ymin:ymax, xmin:xmax] = cv2.resize(img, (xmax-xmin, ymax-ymin))

        # 这里应该添加其他三张图像的拼接逻辑
        # 简化实现，实际项目中需要从数据集随机选择其他图像
        return new_img

    def apply_to_bbox(self, bbox, **params):
        # 调整边界框坐标以适应马赛克增强
        x, y, width, height = bbox
        return [x, y, width, height]

class CutMixAugmentation(DualTransform):
    """CutMix增强实现，将一张图像的区域剪切粘贴到另一张图像"""
    def __init__(self, probability=1.0, alpha=1.0):
        super().__init__(always_apply=False, p=probability)
        self.alpha = alpha

    def apply(self, img, **params):
        # 实现CutMix增强逻辑
        if self.alpha > 0:
            lam = np.random.beta(self.alpha, self.alpha)
        else:
            lam = 1

        h, w = img.shape[:2]
        cut_rat = np.sqrt(1. - lam)
        cut_w = int(w * cut_rat)
        cut_h = int(h * cut_rat)

        # 随机选择剪切区域
        cx = np.random.randint(w)
        cy = np.random.randint(h)

        bbx1 = np.clip(cx - cut_w // 2, 0, w)
        bby1 = np.clip(cy - cut_h // 2, 0, h)
        bbx2 = np.clip(cx + cut_w // 2, 0, w)
        bby2 = np.clip(cy + cut_h // 2, 0, h)

        # 这里应该添加从另一张图像复制区域的逻辑
        # 简化实现，实际项目中需要从数据集随机选择另一张图像
        return img

    def apply_to_bbox(self, bbox, **params):
        # 调整边界框坐标以适应CutMix增强
        x, y, width, height = bbox
        return [x, y, width, height]